# coding=utf-8
from expr_optimal import schedule_optimal
from expr_greedy import schedule_greedy
from expr_swap import schedule_swap
from chiplet import QPU
from dag import Circuit_DAG
import concurrent.futures
import os

max_threads = os.cpu_count()

# print(max_threads)

# 算法类型 bv / qaoa / qft / vqa
# 量子比特数 50 100 261
round_num = 10  # 多次求平均

circuit_type = ["bv", "qaoa", "qft", "vqa"]
qubit_num = [20, 40, 60, 80, 100, 120, 140, 160, 180, 200]
chiplet_type = ["square", "hexagon", "heavy_square", "heavy_hexagon"]

chiplet_size = [
    {"square": 5, "hexagon": 3, "heavy_square": 3, "heavy_hexagon": 2},
    {"square": 5, "hexagon": 3, "heavy_square": 3, "heavy_hexagon": 2},
    {"square": 5, "hexagon": 3, "heavy_square": 3, "heavy_hexagon": 2},
    {"square": 5, "hexagon": 3, "heavy_square": 3, "heavy_hexagon": 2},
    {"square": 5, "hexagon": 3, "heavy_square": 3, "heavy_hexagon": 2},
    {"square": 5, "hexagon": 3, "heavy_square": 3, "heavy_hexagon": 2},
    {"square": 8, "hexagon": 5, "heavy_square": 5, "heavy_hexagon": 3},
    {"square": 8, "hexagon": 5, "heavy_square": 5, "heavy_hexagon": 3},
    {"square": 8, "hexagon": 5, "heavy_square": 5, "heavy_hexagon": 3},
    {"square": 8, "hexagon": 5, "heavy_square": 5, "heavy_hexagon": 3},
]

chiplet_array = [
    {"square": [2, 3], "hexagon": [2, 2],
        "heavy_square": [2, 3], "heavy_hexagon": [2, 2]},
    {"square": [2, 3], "hexagon": [2, 2],
        "heavy_square": [2, 3], "heavy_hexagon": [2, 2]},
    {"square": [2, 3], "hexagon": [2, 2],
        "heavy_square": [2, 3], "heavy_hexagon": [2, 2]},
    {"square": [3, 4], "hexagon": [3, 3],
        "heavy_square": [3, 4], "heavy_hexagon": [3, 3]},
    {"square": [3, 4], "hexagon": [3, 3],
        "heavy_square": [3, 4], "heavy_hexagon": [3, 3]},
    {"square": [3, 4], "hexagon": [3, 3],
        "heavy_square": [3, 4], "heavy_hexagon": [3, 3]},
    {"square": [3, 3], "hexagon": [3, 3],
        "heavy_square": [3, 3], "heavy_hexagon": [3, 3]},
    {"square": [3, 3], "hexagon": [3, 3],
        "heavy_square": [3, 3], "heavy_hexagon": [3, 3]},
    {"square": [3, 3], "hexagon": [3, 3],
        "heavy_square": [3, 3], "heavy_hexagon": [3, 3]},
    {"square": [3, 3], "hexagon": [3, 3],
        "heavy_square": [3, 3], "heavy_hexagon": [3, 3]},
]


def task(num_idx, num_item, circuit_type_item, chiplet_type_item):
    # 这里是原来嵌套循环中的处理逻辑
    # QASM文件
    qasm_file = f"circuits/{circuit_type_item}_{num_item}.qasm"
    # 电路的DAG图
    circuit_dag = Circuit_DAG(qasm_file)
    # Chip内大小
    inner_chip_size = chiplet_size[num_idx][chiplet_type_item]
    # 芯片矩阵大小
    inter_chip_array = chiplet_array[num_idx][chiplet_type_item]
    total_count = {
        "swap": [[], [], []],
        "greedy": [[], [], []],
        "optimal": [[], [], []],
    }
    for _ in range(round_num):
        qpu = QPU(chip_type=chiplet_type_item, inner_chip_size=inner_chip_size, inter_chip_array=inter_chip_array,
                  circuit_dag=circuit_dag, just_one_map=(circuit_type == "vqa" or circuit_type == "qft"))
        for func_idx, func in enumerate([schedule_swap, schedule_greedy, schedule_optimal]):
            max_depth, swap_num, ghz_length = func(qpu)
            arr = ["swap", "greedy", "optimal"]
            total_count[arr[func_idx]][0].append(max_depth)
            total_count[arr[func_idx]][1].append(swap_num)
            total_count[arr[func_idx]][2].append(ghz_length)
    average_count = {
        "swap": [sum(total_count["swap"][0]) / len(total_count["swap"][0]),
                 sum(total_count["swap"][1]) /
                 len(total_count["swap"][1]),
                 sum(total_count["swap"][2]) / len(total_count["swap"][2])],
        "greedy": [sum(total_count["greedy"][0]) / len(total_count["greedy"][0]),
                   sum(total_count["greedy"][1]) /
                   len(total_count["greedy"][1]),
                   sum(total_count["greedy"][2]) / len(total_count["greedy"][2])],
        "optimal": [sum(total_count["optimal"][0]) / len(total_count["optimal"][0]),
                    sum(total_count["optimal"][1]) /
                    len(total_count["optimal"][1]),
                    sum(total_count["optimal"][2]) / len(total_count["optimal"][2])],
    }
    output_file_path = "output.txt"
    with open(output_file_path, 'a') as file:
        file.write(
            f"{num_item} {circuit_type_item} {chiplet_type_item} {inner_chip_size} {inter_chip_array} {average_count}\n")
    print(num_item, circuit_type_item, chiplet_type_item,
          inner_chip_size, inter_chip_array, average_count)


with concurrent.futures.ThreadPoolExecutor(max_workers=max_threads) as executor:
    futures = []
    is_reverse = False
    if (is_reverse):
        # 获取qubit_num的枚举列表并反转
        for num_idx, num_item in reversed(list(enumerate(qubit_num))):
            # "square", "hexagon", "heavy_square", "heavy_hexagon"
            for circuit_type_item in reversed(circuit_type):
                for chiplet_type_item in reversed(chiplet_type):
                    # 提交任务到线程池
                    futures.append(executor.submit(
                        task, num_idx, num_item, circuit_type_item, chiplet_type_item))
    else:
        for num_idx, num_item in enumerate(qubit_num):
            for circuit_type_item in circuit_type:  # "square", "hexagon", "heavy_square", "heavy_hexagon"
                for chiplet_type_item in chiplet_type:
                    # 提交任务到线程池
                    futures.append(executor.submit(
                        task, num_idx, num_item, circuit_type_item, chiplet_type_item))

    # 等待所有任务完成（根据你的需求决定是否需要）
    for future in concurrent.futures.as_completed(futures):
        try:
            future.result()
        except Exception as exc:
            print(f'Generated an exception: {exc}')
